package io.kwdb.sdk;

import io.grpc.ManagedChannel;
import io.grpc.stub.StreamObserver;
import io.kwdb.proto.Ingest;
import io.kwdb.proto.IngestServiceGrpc;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * High-performance streaming client for KWDB using gRPC streaming. Provides client-side streaming
 * for maximum throughput.
 */
public class StreamingClient implements AutoCloseable {
  private static final Logger logger = Logger.getLogger(StreamingClient.class.getName());

  private final ManagedChannel channel;
  private final IngestServiceGrpc.IngestServiceStub asyncStub;
  private StreamObserver<Ingest.WriteRequest> requestObserver;
  private final AtomicLong totalPointsSent = new AtomicLong(0);
  private final AtomicLong totalRowsWritten = new AtomicLong(0);
  private final AtomicReference<Status> lastError = new AtomicReference<>();
  private volatile boolean streamActive = false;
  private CountDownLatch streamLatch;

  /**
   * Create a streaming client with the given channel.
   *
   * @param channel The gRPC channel to use
   */
  public StreamingClient(ManagedChannel channel) {
    this.channel = channel;
    this.asyncStub = IngestServiceGrpc.newStub(channel);
  }

  /**
   * Start a streaming session. Must be called before writing any data.
   *
   * @return CompletableFuture that completes when stream is ready
   */
  public CompletableFuture<Status> startStream() {
    CompletableFuture<Status> future = new CompletableFuture<>();

    if (streamActive) {
      future.complete(Status.error(Status.Code.FAILED_PRECONDITION, "Stream already active"));
      return future;
    }

    streamLatch = new CountDownLatch(1);

    // Create response observer
    StreamObserver<Ingest.WriteResponse> responseObserver =
        new StreamObserver<Ingest.WriteResponse>() {
          @Override
          public void onNext(Ingest.WriteResponse response) {
            totalRowsWritten.addAndGet(response.getRowsWritten());
            logger.log(Level.FINE, "Stream response: {0} rows written", response.getRowsWritten());
          }

          @Override
          public void onError(Throwable t) {
            streamActive = false;
            lastError.set(Status.error(Status.Code.UNKNOWN, t.getMessage()));
            logger.log(Level.WARNING, "Stream error", t);
            streamLatch.countDown();
            future.completeExceptionally(t);
          }

          @Override
          public void onCompleted() {
            streamActive = false;
            logger.log(
                Level.INFO, "Stream completed. Total written: {0} rows", totalRowsWritten.get());
            streamLatch.countDown();
          }
        };

    // Start the stream
    requestObserver = asyncStub.writeStream(responseObserver);
    streamActive = true;

    logger.log(Level.INFO, "Stream started successfully");
    future.complete(Status.ok());

    return future;
  }

  /**
   * Write a batch of points to the stream.
   *
   * @param points List of points to write
   * @return Status of the write operation
   */
  public Status writeBatch(List<Point> points) {
    if (!streamActive) {
      return Status.error(
          Status.Code.FAILED_PRECONDITION, "Stream not active. Call startStream() first");
    }

    if (points == null || points.isEmpty()) {
      return Status.invalidArgument("No points to write");
    }

    try {
      // Convert Points to proto format
      Ingest.WriteRequest.Builder requestBuilder = Ingest.WriteRequest.newBuilder();

      for (Point point : points) {
        Ingest.Point.Builder pointBuilder =
            Ingest.Point.newBuilder()
                .setMeasurement(point.getMeasurement())
                .setTimestampUnixNs(point.getTimestamp());

        // Add tags
        point
            .getTags()
            .forEach(
                (key, value) ->
                    pointBuilder.addTags(
                        Ingest.Tag.newBuilder().setKey(key).setValue(value).build()));

        // Add fields (handle Object type - could be Double, Long, etc.)
        point
            .getFields()
            .forEach(
                (key, value) -> {
                  if (value instanceof Number) {
                    pointBuilder.addFields(
                        Ingest.Field.newBuilder()
                            .setKey(key)
                            .setF64(((Number) value).doubleValue())
                            .build());
                  }
                });

        requestBuilder.addPoints(pointBuilder.build());
      }

      // Send the request
      requestObserver.onNext(requestBuilder.build());
      totalPointsSent.addAndGet(points.size());

      logger.log(Level.FINE, "Sent batch of {0} points", points.size());
      return Status.ok();

    } catch (Exception e) {
      logger.log(Level.WARNING, "Error writing batch", e);
      return Status.error(Status.Code.INTERNAL, e.getMessage());
    }
  }

  /**
   * Write a single point to the stream.
   *
   * @param point Point to write
   * @return Status of the write operation
   */
  public Status writePoint(Point point) {
    return writeBatch(List.of(point));
  }

  /**
   * Finish the stream and wait for completion.
   *
   * @return Status of the stream completion
   */
  public Status finish() {
    return finish(30, TimeUnit.SECONDS);
  }

  /**
   * Finish the stream and wait for completion with timeout.
   *
   * @param timeout Timeout value
   * @param unit Timeout unit
   * @return Status of the stream completion
   */
  public Status finish(long timeout, TimeUnit unit) {
    if (!streamActive) {
      return Status.error(Status.Code.FAILED_PRECONDITION, "Stream not active");
    }

    try {
      // Complete the stream
      requestObserver.onCompleted();

      // Wait for response
      if (!streamLatch.await(timeout, unit)) {
        return Status.error(Status.Code.DEADLINE_EXCEEDED, "Stream finish timeout");
      }

      // Check for errors
      Status error = lastError.get();
      if (error != null) {
        return error;
      }

      logger.log(
          Level.INFO,
          "Stream finished. Sent {0} points, wrote {1} rows",
          new Object[] {totalPointsSent.get(), totalRowsWritten.get()});

      return Status.ok();

    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      return Status.error(Status.Code.ABORTED, "Stream finish interrupted");
    } catch (Exception e) {
      logger.log(Level.WARNING, "Error finishing stream", e);
      return Status.error(Status.Code.INTERNAL, e.getMessage());
    }
  }

  /**
   * Get the total number of points sent in this stream.
   *
   * @return Total points sent
   */
  public long getTotalPointsSent() {
    return totalPointsSent.get();
  }

  /**
   * Get the total number of rows written by the server.
   *
   * @return Total rows written
   */
  public long getTotalRowsWritten() {
    return totalRowsWritten.get();
  }

  /**
   * Check if the stream is currently active.
   *
   * @return True if stream is active, false otherwise
   */
  public boolean isStreamActive() {
    return streamActive;
  }

  @Override
  public void close() {
    if (streamActive) {
      finish(5, TimeUnit.SECONDS);
    }
  }
}
